from base.BaseModel import BaseModel
import pyspark.sql.functions as F
from pyspark.sql.types import StringType


class AgeModel(BaseModel):
    def compute(self, es_df, five_df):
        # 业务数据处理
        es_df = es_df.select("id", F.regexp_replace("birthday", "-", "").alias("birthday"))
        es_df.printSchema()
        es_df.show()
        # 标签规则数据处理
        five_df = five_df.select("id",
                                 F.split("rule", "-")[0].alias("start"),
                                 F.split("rule", "-")[1].alias("end"))
        five_df.printSchema()
        five_df.show()
        # 打标签实现
        new_df = es_df.join(other=five_df,
                            on=es_df['birthday'].between(five_df['start'], five_df['end']),
                            how='left').select(es_df['id'].cast(StringType()).alias("userId"),
                                               five_df['id'].cast(StringType()).alias("tagsId"))
        return new_df

if __name__ == '__main__':
    am = AgeModel(14)
    am.execute()
